# -*- coding: utf-8 -*-

import tushare as ts
import pandas as pd
import matplotlib.pyplot as plt

pd.set_option('display.max_columns',1000)
pd.set_option('display.width', 1000)
pd.set_option('display.max_colwidth',1000)


startDate = '2019-01-01'
endDate = '2019-06-16'

# 时间分3段，用首尾两个时间段来比较收益排名均值的变化
obsPercent = 0.3

# 取收益最高的和最低的各10%，计算收益排名的均值
cmpPercent = 0.1

# 全市场股票信息列表
stockList = ts.get_stock_basics()["name"]
pfList = []

cnt = 0
# print(stockList)

for code, name in stockList.iteritems():
    # print(cnt, code, name)

    # 排除ST和*ST股票
    if name.find("ST") != -1:
        continue

    # 读日K行情数据
    df = ts.get_k_data(code, start=startDate, end=endDate, ktype='D', autype='qfq')

    # 如果有较长的停牌时间，则忽略它
    if len(df) < 50:
        continue

    df.index = df.pop("date")

    # 用于计算收益率的时间窗口宽度
    winSize = int(len(df)  * obsPercent)

    # 第1段时间内的收益率
    profit1 = int((df["close"][winSize-1] / df["close"][0] - 1.0) * 100)

    # 第2段时间内的收益率
    profit2 = int((df["close"][-1] / df["close"][-winSize] - 1.0) * 100)

    # 为便于观察分布图，忽略一些特别大或特别小的值
    if profit1 > 150 or profit2 > 150:
        continue

    # 保存每只股票的2个收益率的值
    pfList.append(dict(code=code, pf1=profit1, pf2=profit2))

    cnt += 1
    # if cnt == 2:
    #     break

# 先按第1段时间内的收益率排名（从高往低）
pfList.sort(key=lambda  x: x["pf1"], reverse=True)
# print(pfList)
# 记录下每只个股在第1段时间内的收益率排名
for idx, doc in enumerate(pfList):
    doc["rank1"] = idx
    # print(idx, doc)

# 用于统计的最高（和最低）收益的股票个数
numCands = int(len(pfList) * cmpPercent)

# 第1段时间内最高收益的股票信息列表
topCands = pfList[:numCands]

# 第1段时间内最低收益的股票信息列表
btmCands = pfList[-numCands:]

# 最高收益股票代码集合
topGroupCodes = set([doc["code"] for doc in topCands])

# 最低收益股票代码集合
btmGroupCodes = set([doc["code"] for doc in btmCands])

# 最高收益股票排名均值
avgTopIdx = round(sum([doc["rank1"] for doc in topCands]) / numCands / len(pfList), 2)

# 最低收益股票排名均值
avgBtmIdx = round(sum([doc["rank1"] for doc in btmCands]) / numCands / len(pfList), 2)

# 所有股票收益排名的均值，等差数列均值（中值），直接计算即可
avgMidIdx = (0 + len(pfList) - 1) / 2 / len(pfList)

# 再按第2段时间内的收益率排名（从高往低
pfList.sort(key=lambda  x: x["pf2"], reverse=True)
# 记录下每只个股在第2段时间内的收益率排名
for idx, doc in enumerate(pfList):
    doc["rank2"] = idx
cmpAvgTopIdx = round(sum([doc["rank2"] for doc in pfList if doc["code"] in topGroupCodes]) / numCands / len(pfList), 2)
cmpAvtBtmIdx = round(sum([doc["rank2"] for doc in pfList if doc["code"] in btmGroupCodes]) / numCands / len(pfList), 2)

print("Top group:", avgTopIdx, "->", cmpAvgTopIdx, "(avg:%.2f)" %avgMidIdx)
print("Btm group:", avgBtmIdx, "->", cmpAvtBtmIdx, "(avg:%.2f" %avgMidIdx)

# profits = pd.DataFrame({"pf1": [doc["pf1"] for doc in pfList]}, {"pf2": [doc["pf2"] for doc in pfList]})
profits = pd.DataFrame({"pf1": [doc["pf1"] for doc in pfList],
"pf2": [doc["pf2"] for doc in pfList]})
# 绘制两段时间内全市场收益率分布的直方图
profits.plot.hist(bins=200, alpha=0.5)
plt.show()
